import requests
import json
import glob
from flask import Flask, request, render_template, send_file
import os
import tempfile

# 百度AI的API Key和Secret Key
API_KEY = "idqZ7gpunT7psq0KsYaRwAuR"
SECRET_KEY = "OHOKDwiTR7Jby4Qb7jWBt2oGzpauvnjs"
BAIDU_AI_URL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_speed"

# Flask应用
app = Flask(__name__)

def get_baidu_access_token():
    """
    使用 AK，SK 生成百度AI的鉴权签名（Access Token）
    :return: access_token，或是None(如果错误)
    """
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
    response = requests.post(url, params=params)
    return response.json().get("access_token")

def read_files(file_path_pattern):
    file_contents = []
    for file_path in glob.glob(file_path_pattern):
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
            file_contents.append(content)
    return file_contents

def check_compliance(law_contents, company_contents):
    results = []
    access_token = get_baidu_access_token()
    for company_content in company_contents:
        for law_content in law_contents:
            SYSTEM_PROMPT = f"""
你将使用中文回答用户的问题，除非用户使用的是其他语言。
下面用户将提供相关法律法规片段和公司的文件，你将参考用户提供的法律法规片段和其他任何可能相关的中国法律，检查这些文件并返回不符合法规的内容
你将回复JSON格式字符串，JSON定义如下：
{{ "result" : [{"content":"<公司文件片段>", "law":"违反《中华人民共和国XXX法》第<>条、第<>条"}]}}
result: 检查的结果list
content: 不符合法规的内容部分，尽可能精简
law: 违反的法律法规，尽可能精简
你将输出完整的plain json字符串，不包括markdown格式等其他内容。

法律片段如下：
{law_content}

公司文件如下：
{company_content}
            """
            headers = {'Content-Type': 'application/json'}
            payload = json.dumps({
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                ]
            })
            url = f"{BAIDU_AI_URL}?access_token={access_token}"
            response = requests.post(url, headers=headers, data=payload)
            try:
                response_json = response.json()
                # 假设百度AI返回的结果在response_json['result']中，具体结构需要根据实际API响应调整
                # 这里假设每个result项都是一个包含"content"和"law"的字典
                results.extend(response_json.get("result", []))
            except json.JSONDecodeError:
                print(f"解析结果出现错误，返回内容为：{response.text}")
    return results

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        if 'law_files' not in request.files or 'company_files' not in request.files:
            return "请上传法律法规文件和公司文件", 400

        law_files = request.files.getlist('law_files')
        company_files = request.files.getlist('company_files')

        with tempfile.TemporaryDirectory() as temp_dir:
            law_file_paths = []
            company_file_paths = []

            for file in law_files:
                file_path = os.path.join(temp_dir, file.filename)
                file.save(file_path)
                law_file_paths.append(file_path)

            for file in company_files:
                file_path = os.path.join(temp_dir, file.filename)
                file.save(file_path)
                company_file_paths.append(file_path)

            law_contents = read_files(os.path.join(temp_dir, '*'))
            company_contents = read_files(os.path.join(temp_dir, '*'))
            compliance_results = check_compliance(law_contents, company_contents)

            data = []
            for item in compliance_results:
                data.append({
                    "content": item.get("content", ""),
                    "law": item.get("law", "")
                })

            with tempfile.NamedTemporaryFile(suffix='.xlsx', delete=False) as tmp_excel:
                import pandas as pd
                df = pd.DataFrame(data)
                df.to_excel(tmp_excel.name, index=False)

            return render_template('result.html', results=compliance_results, file_path=tmp_excel.name)
    return render_template('index.html')

@app.route('/download', methods=['GET'])
def download():
    file_path = request.args.get('file_path')
    return send_file(file_path, as_attachment=True, download_name='results.xlsx')

if __name__ == '__main__':
    app.run(debug=True)